
import numpy as np
from scipy.sparse import coo_matrix
import matplotlib.pyplot as plt
import os
from tqdm import tqdm
import scipy


def construct_X(M, alphas, Np=None):
    if Np is None:
        Np = int(np.ceil(np.sqrt(2) * M))
    if Np % 2 == 0:
        Np = Np + 1  # get an odd number for Np
    print(f"Np={Np}")

    D = M * M
    No = len(alphas)
    N = Np * No

    h = 1
    a0 = - (M - 1) / 2
    b0 = - (M - 1) / 2
    s0 = (Np - 1) / 2 # Set the center of response vector as the middle point of the vector
    ja = np.mgrid[0:D] % M
    jb = np.mgrid[0:D] // M
    C = np.stack((a0 + ja * h, b0 + jb * h)) # move the origin of vectors to the center element of matirx X
    # print(f"C.shape={C.shape}")

    i_indices = []
    j_indices = []
    weights = []

    for alpha_idx in range(len(alphas)):
        alpha = alphas[alpha_idx]
        n = np.array([np.cos(np.deg2rad(alpha)), np.sin(np.deg2rad(alpha))]) # the normal vector of angle alpha
        p = np.dot(n.T, C) + s0

        # j_indices is the column index of X
        # j = ja + M * jb, it's the global index in D dimension
        j_indices.extend(ja + M * jb) 

        # i_indeces is the row index of X
        # add alpha_idx * Np for different alphas
        # p // 1 to get the integer part of p, which is the index on the response vector
        i_indices.extend((p // 1) + alpha_idx * Np) 

        # the weights of X_ij is the decimal part of p
        # need to use 1 - (p % 1) because the smaller the value, the smaller the distance, the larger the impact
        weights.extend(1-(p % 1)) 

        # We suggest that one pixel in image mu affect two bins in the response vector
        # So now we add the second bin
        
        j_indices.extend(ja + M * jb)
        
        # We assume that the pixel only affect the next bin along the index, so (p + 1)
        i_indices.extend(((p + 1) // 1) + alpha_idx * Np)
        
        # For this bin, the weight is 1-(1-(p % 1)) = p % 1
        weights.extend(p % 1)

    X = coo_matrix((weights, (i_indices, j_indices)),
                   shape=(N, D), dtype=np.float32)
    return X


data_path = "./hs_tomography/"

y = np.load(os.path.join(data_path, "y_195.npy"))
alphas = np.load(os.path.join(data_path, "alphas_195.npy"))

# Example usage
print(f"alphas.shape={alphas.shape}")
print(f"y.shape={y.shape}")
alphas = [-33, 1, 42]
X = construct_X(10, alphas)

# Visualize the matrix with the adjusted color scale
fig, axs = plt.subplots(2,1, figsize=(10,8))
im0 = axs[0].imshow(X.toarray(), cmap='viridis')
axs[0].set_title('X matrix')
fig.colorbar(im0, ax=axs[0])

X_example = np.load('hs_tomography/X_example.npy', allow_pickle=True)

im1 = axs[1].imshow(X_example, cmap='viridis')
axs[1].set_title('X_example matrix')
fig.colorbar(im1, ax=axs[1])

# Adjust the spacing between subplots
plt.subplots_adjust(hspace=0.5,top = 0.9)

plt.show()

alphas.shape=(179,)
y.shape=(49225,)
Np=15


y_77 = np.load("hs_tomography/y_77.npy")
alphas_77 = np.load("hs_tomography/alphas_77.npy")
M_77 = 77
X_77 = construct_X(M_77, alphas_77,109)
X_77 = scipy.sparse.csc_matrix(X_77)

print(f"X.shape={X_77.shape}, y.shape={y_77.shape}")

beta_77 = scipy.sparse.linalg.lsqr(X_77,y_77, atol =1e-05, btol = 1e-05)[0]
print(beta_77)

plt.imshow(beta_77.reshape(-1,M_77), cmap='viridis', origin='lower')
plt.colorbar()
plt.title('tomogram')
plt.show()

Np=109
X.shape=(9810, 5929), y.shape=(9810,)
[  2.78814411  14.55389582 222.71980645 ...   0.8825431    1.28509821
   0.59501147]


y_195 = np.load(os.path.join(data_path, "y_195.npy"))
alphas_195 = np.load(os.path.join(data_path, "alphas_195.npy"))
M_195 = 195
X_195 = construct_X(M_195, alphas_195,275)
X_195 = scipy.sparse.csc_matrix(X_195)

print(f"X.shape={X_195.shape}, y.shape={y_195.shape}")

beta_195 = scipy.sparse.linalg.lsqr(X_195,y_195, atol =1e-05, btol = 1e-05)[0]
print(beta_195)

plt.imshow(beta_195.reshape(-1,M_195), cmap='viridis', origin='lower')
plt.colorbar()
plt.title('tomogram')
plt.show()

Np=275
X.shape=(49225, 38025), y.shape=(49225,)
[63.03236529 64.3488428  63.39724164 ... 64.60526625 63.10795307
 61.771356  ]


print(f"X.shape={X_195.shape}, y.shape={y_195.shape}")

beta_195 = scipy.sparse.linalg.lsqr(X_195,y_195, atol =1e-08, btol = 1e-08)[0]
print(beta_195)

plt.imshow(beta_195.reshape(-1,M_195), cmap='viridis', origin='lower')
plt.colorbar()
plt.title('tomogram')
plt.show()

X.shape=(49225, 38025), y.shape=(49225,)
[62.98468007 63.95911918 64.68461242 ... 64.65732334 63.06304968
 61.97625969]


y_195 = np.load("hs_tomography/y_195.npy")
alphas_195 = np.load("hs_tomography/alphas_195.npy")
M_195 = 195
Np=275

max_stride = 10
for i in tqdm(range(max_stride)): # test for all strides from 1 to max_stride
    alphas_test = alphas_195[::i+1]
    X_test = construct_X(M_195, alphas_test ,Np) # sample from angles
    X_test = scipy.sparse.csc_matrix(X_test)

    # take #Np rows per angle
    # Create y_test by selecting consecutive rows with the given stride
    y_test = []
    for i2 in range(len(alphas_test)):
        i3 = i2 * (i + 1)
        y_test.extend(y_195[i3 * Np : (i3+1) * Np])
    
    y_test = np.array(y_test)

    print(f"X.shape={X_test.shape}, y.shape={y_test.shape}")

    beta_195 = scipy.sparse.linalg.lsqr(X_test,y_test, atol =1e-05, btol = 1e-05)[0]

    plt.figure()
    plt.imshow(beta_195.reshape(-1,M_195), cmap='viridis', origin='lower')
    plt.title(f'#alpha={len(alphas_test)}')
    plt.colorbar()
    plt.show()

# Adjust the spacing between subplots
# plt.subplots_adjust(hspace=0.5,top = 0.9)

# plt.show()

  0%|          | 0/10 [00:00<?, ?it/s]

Np=275
X.shape=(49225, 38025), y.shape=(49225,)

 10%|█         | 1/10 [00:22<03:18, 22.04s/it]

Np=275
X.shape=(24750, 38025), y.shape=(24750,)

 20%|██        | 2/10 [00:32<02:00, 15.02s/it]

Np=275
X.shape=(16500, 38025), y.shape=(16500,)

 30%|███       | 3/10 [00:38<01:16, 10.98s/it]

Np=275
X.shape=(12375, 38025), y.shape=(12375,)

 40%|████      | 4/10 [00:42<00:50,  8.47s/it]

Np=275
X.shape=(9900, 38025), y.shape=(9900,)

Exercise 4¶

1. Constructing the matrix X¶

2. Recovering the image¶

3 Minimizing the radiation dose¶